import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os


class MultiModalDataset(Dataset):
    def __init__(self, txt_file, transform=None):
        """
        txt_file: 包含图片路径和标签的txt文件路径
        transform: 预处理操作（例如：图像缩放、归一化等）
        """
        self.data = []
        self.transform = transform

        # 读取txt文件并解析图片路径和标签
        with open(txt_file, 'r') as f:
            lines = f.readlines()
            for line in lines:
                paths_and_label = line.strip().split(',')
                DIFF_image_path = paths_and_label[0].strip()  # motley 图像路径
                WNB_image_path = paths_and_label[1].strip()  # cyan 图像路径
                label = int(paths_and_label[2].strip())  # 标签（0, 1, 2）

                # 将路径和标签保存为元组
                self.data.append((DIFF_image_path, WNB_image_path, label))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        DIFF_image_path, WNB_image_path, label = self.data[idx]

        # 读取图片
        DIFF_image = Image.open(DIFF_image_path).convert("RGB")
        WNB_image = Image.open(WNB_image_path).convert("RGB")

        if self.transform:
            DIFF_image = self.transform(DIFF_image)
            WNB_image = self.transform(WNB_image)

        return DIFF_image, WNB_image, label


# 定义图像预处理（例如：缩放、归一化等）
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 假设你使用ResNet50并且输入尺寸是224x224
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # ImageNet的标准化值
])

# 使用该数据集
# txt_file = 'data_Diff_Wnb_train_new.txt'  # 你的txt文件路径
# dataset = MultiModalDataset(txt_file=txt_file, transform=transform)
# dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

# 测试
# for DIFF_image, WNB_image, labels in dataloader:
#     print(DIFF_image.shape)  # 打印motley图像的形状
#     print(WNB_image.shape)  # 打印cyan图像的形状
#     break

# 扩展写入datas.txt文件的部分，增加新的图像路径
# DIFF_dir = r'E:\Research topic\DIFF_WNB\DIFF'
DIFF_dir = r'D:\competition\2025tongji\DIFF_WNB_hospital\val\DIFF'
# WNB_dir = r'E:\Research topic\DIFF_WNB\WNB'
WNB_dir = r'D:\competition\2025tongji\DIFF_WNB_hospital\val\WNB'
subfolders = ['All_high', 'All_relapse', 'All_correlation', 'Nornmal']
output_file = 'data_Diff_Wnb_val2.txt'  # 修改为datas.txt

# 用于存储每个子文件夹中的图片数量
image_count_per_category = {}

with open(output_file, 'w') as f:
    # 遍历每个子文件夹和标签
    for label, subfolder in enumerate(subfolders):
        # 构建DIFF和WNB图像的路径

        DIFF_subfolder_path = os.path.join(DIFF_dir, subfolder)
        WNB_subfolder_path = os.path.join(WNB_dir, subfolder)

        # 初始化图片数量
        image_count = 0
        # 遍历DIFF图像子文件夹中的所有文件
        for diff_image in os.listdir(DIFF_subfolder_path):
            if diff_image.endswith('.png'):  # 假设文件格式为png

                DIFF_image_path = os.path.join(DIFF_subfolder_path, diff_image)
                WNB_image_path = os.path.join(WNB_subfolder_path, diff_image)  # 假设两个文件夹中的文件名一致

                # 写入文件，格式：DIFF_image_path, WNB_image_path, label
                f.write(f"{DIFF_image_path},{WNB_image_path},{label}\n")
                # 增加图片数量
                image_count += 1

        # 记录每个类别中的图片数量
        image_count_per_category[subfolder] = image_count
# 打印每个类别中的图片数量
print("Image count per category:")
for category, count in image_count_per_category.items():
    print(f"{category}: {count} images")